package com.example.demo.utils;

import com.dampcake.bencode.Bencode;
import com.dampcake.bencode.Type;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;


public class BencodeUtil {

    private static Bencode bencodeUTF8 = new Bencode(StandardCharsets.UTF_8);
    private static Bencode bencodeInfoHash = new Bencode(StandardCharsets.ISO_8859_1);
    // ----------------- 编码部分 -----------------

    public static byte[] encode(Object data) throws IOException {
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        encodeValue(data, output);
        return output.toByteArray();
    }

    public static String encodeUTF8(Map<String, Object> map) {
        byte[] encodedBytes = bencodeUTF8.encode(map);
        return new String(encodedBytes, StandardCharsets.UTF_8);
    }

    public static String encodeISO(Map<String, Object> map) {
        byte[] encodedBytes = bencodeInfoHash.encode(map);
        return new String(encodedBytes, StandardCharsets.ISO_8859_1);
    }

    public static byte[] encodeISOBytes(Map<String, Object> map) {
        byte[] encodedBytes =bencodeInfoHash.encode(map);
        return encodedBytes;
    }

    public static Map<String, Object> decodeISO(byte[] torrentBytes) throws IOException {
        return bencodeInfoHash.decode(torrentBytes, Type.DICTIONARY);
    }
    private static void encodeValue(Object data, ByteArrayOutputStream output) throws IOException {
        if (data instanceof String) {
            encodeString((String) data, output);
        } else if (data instanceof byte[]) {
            encodeBytes((byte[]) data, output);
        } else if (data instanceof Number) {
            encodeNumber((Number) data, output);
        } else if (data instanceof List) {
            encodeList((List<?>) data, output);
        } else if (data instanceof Map) {
            encodeMap((Map<String, ?>) data, output);
        } else {
            throw new IllegalArgumentException("Unsupported data type: " + data.getClass());
        }
    }

    // 新增流式编码方法
    public static void encodeToStream(Object data, ByteArrayOutputStream output) throws IOException {
        encodeValue(data, output);
    }

    private static void encodeString(String str, ByteArrayOutputStream output) throws IOException {
        byte[] bytes = str.getBytes(StandardCharsets.UTF_8);
        encodeBytes(bytes, output);
    }

    private static void encodeBytes(byte[] bytes, ByteArrayOutputStream output) throws IOException {
        output.write(Integer.toString(bytes.length).getBytes());
        output.write(':');
        output.write(bytes, 0, bytes.length);
    }

    private static void encodeNumber(Number number, ByteArrayOutputStream output) throws IOException {
        output.write('i');
        output.write(number.toString().getBytes(StandardCharsets.UTF_8));
        output.write('e');
    }

    private static void encodeList(List<?> list, ByteArrayOutputStream output) throws IOException {
        output.write('l');
        for (Object item : list) {
            encodeValue(item, output);
        }
        output.write('e');
    }

    private static void encodeMap(Map<String, ?> map, ByteArrayOutputStream output) throws IOException {
        // 对键进行排序
        TreeMap<String, Object> sortedMap = new TreeMap<>(map);
        output.write('d');
        for (Map.Entry<String, ?> entry : sortedMap.entrySet()) {
            encodeString(entry.getKey(), output);
            encodeValue(entry.getValue(), output);
        }
        output.write('e');
    }

    // ----------------- 解码部分 -----------------

    public static Object decode(byte[] data) throws IOException {
        BencodeParser parser = new BencodeParser(data);
        return parser.parse();
    }

    private static class BencodeParser {
        private final byte[] data;
        private int position;

        public BencodeParser(byte[] data) {
            this.data = data;
            this.position = 0;
        }

        public Object parse() throws IOException {
            byte current = nextByte();
            if (current == 'i') {
                return parseNumber();
            } else if (Character.isDigit(current)) {
                return parseBytes();
            } else if (current == 'l') {
                return parseList();
            } else if (current == 'd') {
                return parseMap();
            } else {
                throw new IOException("Invalid bencode format at position " + (position-1));
            }
        }

        private Long parseNumber() throws IOException {
            int start = position;
            while (nextByte() != 'e') {
                // 继续读取直到遇到'e'
            }
            String numStr = new String(data, start, position - start - 1, StandardCharsets.UTF_8);
            return Long.parseLong(numStr);
        }

        private byte[] parseBytes() throws IOException {
            int colonIndex = position;
            while (data[colonIndex] != ':') {
                colonIndex++;
            }
            String lenStr = new String(data, position-1, colonIndex - position +1, StandardCharsets.UTF_8);
            int length = Integer.parseInt(lenStr);
            position = colonIndex + 1;
            byte[] bytes = Arrays.copyOfRange(data, position, position + length);
            position += length;
            return bytes;
        }

        private List<Object> parseList() throws IOException {
            List<Object> list = new ArrayList<>();
            while (data[position] != 'e') {
                list.add(parse());
            }
            position++; // 跳过'e'
            return list;
        }

        private Map<String, Object> parseMap() throws IOException {
            Map<String, Object> map = new LinkedHashMap<>();
            while (data[position] != 'e') {
                byte[] keyBytes = parseBytes();
                String key = new String(keyBytes, StandardCharsets.UTF_8);
                Object value = parse();
                map.put(key, value);
            }
            position++; // 跳过'e'
            return map;
        }

        private byte nextByte() {
            if (position >= data.length) {
                throw new IndexOutOfBoundsException("Unexpected end of data");
            }
            return data[position++];
        }
    }
}